import os
import json
import random
from typing import Dict, Any, List, Tuple, Union, Optional

from PIL import Image, ImageDraw, ImageFont


def _ensure_dir(p: str):
    os.makedirs(p, exist_ok=True)


def _denorm_bbox(b: List[float], W: int, H: int) -> Tuple[int, int, int, int]:
    x1, y1, x2, y2 = b
    x1 = max(0.0, min(1.0, float(x1)))
    y1 = max(0.0, min(1.0, float(y1)))
    x2 = max(0.0, min(1.0, float(x2)))
    y2 = max(0.0, min(1.0, float(y2)))
    x1, x2 = (x1, x2) if x1 <= x2 else (x2, x1)
    y1, y2 = (y1, y2) if y1 <= y2 else (y2, y1)
    return (int(round(x1 * W)), int(round(y1 * H)), int(round(x2 * W)), int(round(y2 * H)))


def _denorm_point(p: List[float], W: int, H: int) -> Tuple[int, int]:
    x, y = p
    x = max(0.0, min(1.0, float(x)))
    y = max(0.0, min(1.0, float(y)))
    return (int(round(x * W)), int(round(y * H)))


def _measure_text(draw: ImageDraw.ImageDraw, text: str, font: ImageFont.ImageFont) -> Tuple[int, int]:
    # Pillow >= 8: draw.textbbox; 再退 font.getbbox / getsize
    try:
        bbox = draw.textbbox((0, 0), text, font=font)
        return int(bbox[2] - bbox[0]), int(bbox[3] - bbox[1])
    except Exception:
        pass
    if hasattr(font, "getbbox"):
        try:
            l, t, r, b = font.getbbox(text)
            return int(r - l), int(b - t)
        except Exception:
            pass
    if hasattr(font, "getsize"):
        try:
            return font.getsize(text)
        except Exception:
            pass
    # 仅用于文字尺寸估计，主流程仍 fail-fast
    return max(8, int(len(text) * 7)), 12


def _is_number(x: Any) -> bool:
    try:
        float(x)
        return True
    except Exception:
        return False


def _normalize_points_field(
    pts_field: Any,
    image_name: str,
    cls_name: str
) -> List[List[float]]:
    """
    接受以下合法形态，并统一返回 [[x,y], ...]:
      - [x, y]
      - [[x, y], [x, y], ...]
      - [(x, y), ...] / (x, y)
      - {"x":..., "y":...} 或 [{"x":...,"y":...}, ...]
    其他形态直接抛错（带上下文）。
    """
    # dict {"x":..,"y":..} → 单点
    if isinstance(pts_field, dict) and "x" in pts_field and "y" in pts_field:
        x, y = pts_field["x"], pts_field["y"]
        if not (_is_number(x) and _is_number(y)):
            raise ValueError(f"[{image_name}] location '{cls_name}' bad dict point: {pts_field}")
        return [[float(x), float(y)]]

    # list/tuple
    if isinstance(pts_field, (list, tuple)):
        # 形如 [x,y]
        if len(pts_field) == 2 and _is_number(pts_field[0]) and _is_number(pts_field[1]):
            return [[float(pts_field[0]), float(pts_field[1])]]
        # 形如 [[x,y], [x,y], ...] 或 [{"x":..,"y":..}, ...]
        normalized: List[List[float]] = []
        for idx, el in enumerate(pts_field):
            if isinstance(el, dict) and "x" in el and "y" in el:
                x, y = el["x"], el["y"]
                if not (_is_number(x) and _is_number(y)):
                    raise ValueError(f"[{image_name}] location '{cls_name}' bad dict point at idx={idx}: {el}")
                normalized.append([float(x), float(y)])
            elif isinstance(el, (list, tuple)) and len(el) == 2 and _is_number(el[0]) and _is_number(el[1]):
                normalized.append([float(el[0]), float(el[1])])
            else:
                raise ValueError(
                    f"[{image_name}] location '{cls_name}' expects [x,y] or list of [x,y]; "
                    f"got invalid element at idx={idx}: {el}"
                )
        return normalized

    raise ValueError(
        f"[{image_name}] location '{cls_name}' expects [x,y] or [[x,y],...]/dict form; got: {type(pts_field).__name__}"
    )


def _normalize_bboxes_field(
    boxes_field: Any,
    image_name: str,
    cls_name: str
) -> List[List[float]]:
    """
    接受以下合法形态，并统一返回 [[x1,y1,x2,y2], ...]:
      - [x1,y1,x2,y2]
      - [[x1,y1,x2,y2], ...]
      - 其他形态直接抛错（带上下文）
    """
    if isinstance(boxes_field, (list, tuple)):
        # 单框
        if len(boxes_field) == 4 and all(_is_number(v) for v in boxes_field):
            x1, y1, x2, y2 = map(float, boxes_field)
            return [[x1, y1, x2, y2]]
        # 多框
        normalized: List[List[float]] = []
        for idx, el in enumerate(boxes_field):
            if isinstance(el, (list, tuple)) and len(el) == 4 and all(_is_number(v) for v in el):
                x1, y1, x2, y2 = map(float, el)
                normalized.append([x1, y1, x2, y2])
            else:
                raise ValueError(
                    f"[{image_name}] detection '{cls_name}' expects [x1,y1,x2,y2] or list of them; "
                    f"got invalid element at idx={idx}: {el}"
                )
        return normalized

    raise ValueError(
        f"[{image_name}] detection '{cls_name}' expects [x1,y1,x2,y2] or [[...],...]; got: {type(boxes_field).__name__}"
    )


def _draw_annotations(img_path: str,
                      item: Dict[str, Any],
                      out_path: str,
                      point_radius_px: Optional[int] = None) -> None:
    img = Image.open(img_path).convert("RGB")
    W, H = img.size
    draw = ImageDraw.Draw(img)

    if point_radius_px is None:
        point_radius_px = max(3, int(round(min(W, H) * 0.006)))

    # 明确失败就报错（不 try/except）
    font = ImageFont.load_default()

    COLORS = [
        (255, 0, 0), (0, 180, 0), (0, 0, 255),
        (255, 140, 0), (200, 0, 200), (0, 160, 160)
    ]

    # --- 检测框 ---
    det = item.get("detection") or {}
    for i, (cls_name, boxes_field) in enumerate(det.items()):
        color = COLORS[i % len(COLORS)]
        boxes = _normalize_bboxes_field(boxes_field, item.get("image_name", "<?>"), cls_name)
        for b in boxes:
            x1, y1, x2, y2 = _denorm_bbox(b, W, H)
            draw.rectangle([x1, y1, x2, y2], outline=color, width=max(2, int(round(min(W, H) * 0.002))))
            label = str(cls_name)
            tw, th = _measure_text(draw, label, font)
            pad = 2
            y_top = max(0, y1 - th - 2 * pad)
            draw.rectangle([x1, y_top, x1 + tw + 2 * pad, y_top + th + 2 * pad], fill=color)
            draw.text((x1 + pad, y_top + pad), label, fill=(255, 255, 255), font=font)

    # --- 定位点 ---
    loc = item.get("location") or {}
    for i, (cls_name, pts_field) in enumerate(loc.items()):
        color = COLORS[(i + 3) % len(COLORS)]
        pts = _normalize_points_field(pts_field, item.get("image_name", "<?>"), cls_name)
        for p in pts:
            cx, cy = _denorm_point(p, W, H)
            r = point_radius_px
            draw.ellipse([cx - r, cy - r, cx + r, cy + r], outline=color, width=max(2, int(round(min(W, H) * 0.002))))
            draw.line([cx - r, cy, cx + r, cy], fill=color, width=1)
            draw.line([cx, cy - r, cx, cy + r], fill=color, width=1)
            label = str(cls_name)
            tw, th = _measure_text(draw, label, font)
            offset = r + 3
            x0 = min(W - tw - 4, max(0, cx + offset))
            y0 = max(0, cy - th - 2)
            draw.rectangle([x0, y0, x0 + tw + 4, y0 + th + 2], fill=color)
            draw.text((x0 + 2, y0), label, fill=(255, 255, 255), font=font)

    img.save(out_path)


def visualize_annotations(annotations: Union[str, Dict[str, Any]],
                          save_dir: str,
                          visual_number: int = 50,
                          seed: int = 42) -> Dict[str, Any]:
    _ensure_dir(save_dir)

    if isinstance(annotations, str):
        with open(annotations, "r", encoding="utf-8") as f:
            ann = json.load(f)
    else:
        ann = annotations

    items: List[Tuple[str, Dict[str, Any]]] = list(ann.items())
    N = len(items)
    if N == 0:
        return {"total": 0, "visualized": 0, "save_dir": save_dir}

    rng = random.Random(seed)
    rng.shuffle(items)
    sample = items[:min(visual_number, N)]

    report = []
    for idx, (img_name, item) in enumerate(sample, start=1):
        img_path = item.get("image_path")
        if not img_path or not os.path.exists(img_path):
            report.append({"image_name": img_name, "status": "skip_missing_image", "image_path": img_path})
            continue

        folder = os.path.join(save_dir, f"{idx:04d}")
        _ensure_dir(folder)

        info_path = os.path.join(folder, "info.txt")
        with open(info_path, "w", encoding="utf-8") as f:
            f.write(json.dumps(item, ensure_ascii=False, indent=2))

        viz_path = os.path.join(folder, "viz.png")
        # 让异常自然抛出（Fail Fast）
        _draw_annotations(img_path, item, viz_path)
        report.append({"image_name": img_name, "status": "ok", "folder": folder})

    with open(os.path.join(save_dir, "summary.json"), "w", encoding="utf-8") as f:
        json.dump(report, f, ensure_ascii=False, indent=2)

    return {"total": N, "visualized": len(sample), "save_dir": save_dir}


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Visualize enhanced coordinate annotations.")
    parser.add_argument("--annotations", type=str, default='./experiments/interim/coor_qa/coor_qa_annotations.json', help="Path to annotations.json")
    parser.add_argument("--save_dir", type=str, default='./tmp/coo_check', help="Directory to save visualization folders")
    parser.add_argument("--visual_number", type=int, default=50, help="Number of samples to visualize")
    parser.add_argument("--seed", type=int, default=42, help="Random seed")
    args = parser.parse_args()
    summary = visualize_annotations(annotations=args.annotations, save_dir=args.save_dir,
                                    visual_number=args.visual_number, seed=args.seed)
    print(summary)
